from flask import Flask, render_template, request, jsonify, redirect, url_for
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
import os
from flask_sqlalchemy import SQLAlchemy
from PIL import Image
import io
from datetime import datetime

# 创建Flask应用实例
app = Flask(__name__)

# 设置上传文件的存储路径
app.config['UPLOAD_FOLDER'] = 'static/uploads'

# 配置数据库连接
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///predictions.db'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
db = SQLAlchemy(app)

class Prediction(db.Model):
    id = db.Column(db.Integer, primary_key=True)  # 主键
    image_url = db.Column(db.String(255), nullable=False)  # 图片路径
    predicted_class = db.Column(db.String(50), nullable=False)  # 预测类别
    confidence = db.Column(db.Float, nullable=False)  # 预测置信度
    timestamp = db.Column(db.DateTime, default=datetime.now)  # 使用本地时间

# 初始化数据库，创建所有表
with app.app_context():
    db.create_all()

# 定义类别与Emoji的映射关系
emoji_map = {
    '飞机': '✈️',
    '汽车': '🚗',
    '鸟': '🐦',
    '猫': '🐱',
    '鹿': '🦌',
    '狗': '🐶',
    '青蛙': '🐸',
    '马': '🐎',
    '船': '🚢',
    '卡车': '🚚'
}

# 延迟加载模型
model = None

def load_model_once():
    global model
    if model is None:
        model = load_model('D:\人工智能生成内容\photo-classifier\model\improved_classifier_model.h5')

# 压缩图片
def compress_image(file, max_size=(800, 800), quality=85):
    """压缩图片"""
    img = Image.open(file)
    img.thumbnail(max_size)  # 调整图片大小
    img_io = io.BytesIO()
    img.save(img_io, format='JPEG', quality=quality)  # 保存为JPEG格式并压缩
    img_io.seek(0)
    return img_io

# 定义预测函数，输入图片路径，返回预测类别和置信度
def predict_image(img_path):
    load_model_once()  # 确保模型已加载
    img = image.load_img(img_path, target_size=(32, 32))
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0) / 255.0
    predictions = model.predict(img_array)
    predicted_class = np.argmax(predictions[0])
    confidence = float(np.max(predictions[0]))
    return list(emoji_map.keys())[predicted_class], confidence

# 定义主页路由，返回主页模板
@app.route('/', methods=['GET'])
def index():
    return render_template('index.html')

# 定义预测路由，处理POST请求
@app.route('/predict', methods=['POST'])
def predict():
    if 'files[]' not in request.files:
        return jsonify({'error': '未选择文件'})
    
    files = request.files.getlist('files[]')
    results = []
    
    for file in files:
        if file.filename == '':
            continue
        
        # 压缩图片
        compressed_file = compress_image(file)
        file_path = os.path.join(app.config['UPLOAD_FOLDER'], file.filename)
        with open(file_path, 'wb') as f:
            f.write(compressed_file.getvalue())
        
        class_name, confidence = predict_image(file_path)
        emoji = emoji_map.get(class_name, '')
        
        new_prediction = Prediction(
            image_url=file_path,
            predicted_class=class_name,
            confidence=confidence
        )
        db.session.add(new_prediction)
        db.session.commit()
        
        results.append({
            'class': class_name,
            'confidence': round(confidence * 100, 2),
            'image_url': file_path,
            'emoji': emoji
        })
    
    return jsonify(results)

# 定义历史记录路由，返回历史记录页面
@app.route('/history', methods=['GET'])
def history():
    category = request.args.get('category')
    page = request.args.get('page', 1, type=int)
    per_page = 10  # 每页显示10条记录

    if category:
        predictions = Prediction.query.filter_by(predicted_class=category).paginate(page=page, per_page=per_page)
    else:
        predictions = Prediction.query.paginate(page=page, per_page=per_page)

    return render_template('history.html', predictions=predictions, category=category)

# 定义清除历史记录路由，处理POST请求
@app.route('/clear_history', methods=['POST'])
def clear_history():
    Prediction.query.delete()
    db.session.commit()
    return redirect(url_for('history'))

# 删除单条历史记录路由
@app.route('/delete/<int:id>', methods=['POST'])
def delete_single(id):
    prediction = Prediction.query.get_or_404(id)
    db.session.delete(prediction)
    db.session.commit()
    return redirect(url_for('history'))

# 批量删除历史记录路由
@app.route('/bulk_delete', methods=['POST'])
def bulk_delete():
    ids = request.form.getlist('prediction_ids')
    for id in ids:
        prediction = Prediction.query.get(int(id))
        if prediction:
            db.session.delete(prediction)
    db.session.commit()
    return redirect(url_for('history'))

# 启动Flask应用
if __name__ == '__main__':
    app.run(debug=True)